package org.papillion39.sechelper.auth;

import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Enumeration;
import java.util.Hashtable;
import java.util.Properties;
import java.util.logging.Logger;

import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.naming.Binding;
import javax.naming.Context;
import javax.naming.InitialContext;
import javax.naming.Name;
import javax.naming.NameClassPair;
import javax.naming.NameParser;
import javax.naming.NamingEnumeration;
import javax.naming.NamingException;
import javax.sql.DataSource;

import org.jvnet.hk2.annotations.Service;
import com.sun.appserv.security.AppservRealm;
import com.sun.enterprise.security.auth.realm.BadRealmException;
import com.sun.enterprise.security.auth.realm.InvalidOperationException;
import com.sun.enterprise.security.auth.realm.NoSuchRealmException;
import com.sun.enterprise.security.auth.realm.NoSuchUserException;
import java.math.BigInteger;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;
import java.util.Vector;
import java.util.logging.Level;

/**
 *
 * @author Jonathan Paz <jonathan.paz@ottr.com>
 */
@Service
public class JDBCSaltedDigestRealm extends AppservRealm {
    private static final Logger LOG = Logger.getLogger(JDBCSaltedDigestRealm.class.getName());
    
    private static final String AUTH_TYPE = "jdbcSaltedDigest";
    private static final String DATASOURCE_JNDI_PARAM = "datasource-jndi";
    private static final String USER_TABLE_PARAM = "user-table";
    private static final String GROUP_TABLE_PARAM = "group-table";
    private static final String USER_GROUP_TABLE_PARAM = "user-group-table";
    private static final String USER_NAME_COLUMN_PARAM = "user-name-column";
    private static final String USER_TABLE_ID_PARAM = "user-table_id";
    private static final String GROUP_TABLE_ID_PARAM = "group-table_id";
    private static final String USER_GROUP_TABLE_USER_ID_PARAM = "user-group-table-user-id";
    private static final String USER_GROUP_TABLE_GROUP_ID_PARAM = "user-group-table-group-id";
    private static final String PASSWORD_COLUMN_PARAM = "password-column";
    private static final String SALT_COLUMN_PARAM = "salt-column";
    private static final String ROUNDS_PARAM = "rounds";
    private static final String GROUP_NAME_COLUMN_PARAM = "group-name-column";
    private static final String ALGORITHM_PARAM = "algorithm";
    private static final String DB_USER_PARAM = "db-user";
    private static final String DB_PASSWORD_PARAM = "db-password";
    private static final String DEFAULT_SALT_LENGTH_PARAM = "default-salt-length";
    private static final String DEFAULT_KEY_LENGTH_PARAM = "default-key-length";
    private static final int DEFAULT_SALT_LENGTH = 256;
    private static final String DEFAULT_ALGORITHM = "PBKDF2WithHmacSHA1";
    private static final int DEFAULT_ROUNDS = 4096;
    private static final int DEFAULT_KEY_LENGTH = 256;

    private static final String BAD_VALUE_MSG = "Illegal value for %s: %s";
    private SecretKeyFactory fac;

    @Override
    protected void init(Properties props) throws BadRealmException, NoSuchRealmException {
        super.init(props); //To change body of generated methods, choose Tools | Templates.
        try {
            setNonNullProperty(props, JAAS_CONTEXT_PARAM);
            setNonNullProperty(props, DATASOURCE_JNDI_PARAM);
            setNonNullProperty(props, USER_TABLE_PARAM);
            setNonNullProperty(props, GROUP_TABLE_PARAM);
            setNonNullProperty(props, USER_NAME_COLUMN_PARAM);
            setNonNullProperty(props, PASSWORD_COLUMN_PARAM);
            setNonNullProperty(props, SALT_COLUMN_PARAM);
            setNonNullProperty(props, GROUP_NAME_COLUMN_PARAM);
            setNullSafeProperty(props, USER_GROUP_TABLE_PARAM);
            setNullSafeProperty(props, USER_TABLE_ID_PARAM);
            setNullSafeProperty(props, GROUP_TABLE_ID_PARAM);
            setNullSafeProperty(props, USER_GROUP_TABLE_USER_ID_PARAM);
            setNullSafeProperty(props, USER_GROUP_TABLE_GROUP_ID_PARAM);

            String algorithm = setDefaultProperty(props, ALGORITHM_PARAM, DEFAULT_ALGORITHM);
            try {
                fac = SecretKeyFactory.getInstance(algorithm);
            } catch (NoSuchAlgorithmException e) {
                BadRealmException ee = new BadRealmException(String.format("No such algorithm exists: %s", algorithm));
                LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "init", e);
                throw ee;
            }

            setDefaultNumberProperty(props, ROUNDS_PARAM, DEFAULT_ROUNDS);
            setDefaultNumberProperty(props, DEFAULT_SALT_LENGTH_PARAM, DEFAULT_SALT_LENGTH);
            setDefaultNumberProperty(props, DEFAULT_KEY_LENGTH_PARAM, DEFAULT_KEY_LENGTH);

            String dbUser = props.getProperty(DB_USER_PARAM);
            String dbPassword = props.getProperty(DB_PASSWORD_PARAM);
            if (dbUser != null && dbPassword != null) {
                setProperty(DB_USER_PARAM, dbUser);
                setProperty(DB_PASSWORD_PARAM, dbPassword);
            }
        } catch (BadRealmException | RuntimeException e) {
            LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "init", e);
            throw e;
        }
    }

    @Override
    public synchronized void setProperty(String name, String value) {
        if (LOG.isLoggable(Level.FINER)) {
            LOG.log(Level.FINER, "Setting property {0} = {1}", new Object[]{name, value});
        }
        super.setProperty(name, value);
    }

    private String setNonNullProperty(Properties props, String prop) throws BadRealmException {
        String value = props.getProperty(prop);
        if (value == null) {
            BadRealmException e = new BadRealmException(String.format("Parameter not set: %s",prop));
            LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "setNonNullProperty", e);
            throw e;
        }
        setProperty(prop,value);
        return value;
    }
    
    private String setNullSafeProperty(Properties props, String prop) {
        String value = props.getProperty(prop);
        if (value != null) {
            setProperty(prop,value);
        }
        return value;
    }

    private String setDefaultProperty(Properties props, String prop, String def) {
        String value = props.getProperty(prop);
        if (value == null) {
            value = def;
        }
        setProperty(prop,value);
        return value;
    }

    private String setDefaultNumberProperty(Properties props, String prop, int def) throws BadRealmException {
        String value = props.getProperty(prop);
        if (value == null) {
            value = Integer.toString(def);
        } else {
            try {
                Integer.parseInt(value);
            } catch (NumberFormatException e) {
                BadRealmException ee = new BadRealmException(String.format(BAD_VALUE_MSG, prop, value),e);
                LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "setDefaultNumberProperty", ee);
                throw ee;
            }
        }
        setProperty(prop,value);
        return value;
    }
    
    @Override
    public String getAuthType() {
        return AUTH_TYPE;
    }

    private void determineUserGroupsSchema() throws BadRealmException {
        String userTableId = getProperty(USER_TABLE_ID_PARAM);
        String groupTableId = getProperty(GROUP_TABLE_ID_PARAM);
        String userGroupTableUserId = getProperty(USER_GROUP_TABLE_USER_ID_PARAM);
        String userGroupTableGroupId = getProperty(USER_GROUP_TABLE_GROUP_ID_PARAM);
        if (userTableId == null || groupTableId == null || userGroupTableUserId == null || userGroupTableGroupId == null) {
            String userTable = getProperty(USER_TABLE_PARAM);
            String groupTable = getProperty(GROUP_TABLE_PARAM);
            String userGroupTable = getProperty(USER_GROUP_TABLE_PARAM);
            LOG.finer("One of the relationship tables is not filled in");
            try (Connection conn = getConnection()) {
                DatabaseMetaData md = conn.getMetaData();
                if (LOG.isLoggable(Level.FINER)) {
                    LOG.log(Level.FINER, "Getting references between {0} and {1}", new Object[]{userTable, userGroupTable});
                }
                try (ResultSet rs = md.getCrossReference(null, null, userTable, null, null, userGroupTable)) {
                    if (rs.next()) {
                        userTableId = rs.getString("PKCOLUMN_NAME");
                        setProperty(USER_TABLE_ID_PARAM,userTableId);
                        userGroupTableUserId = rs.getString("FKCOLUMN_NAME");
                        setProperty(USER_GROUP_TABLE_USER_ID_PARAM,userGroupTableUserId);
                    }
                }
                if (LOG.isLoggable(Level.FINER)) {
                    LOG.log(Level.FINER, "Getting references between {0} and {1}", new Object[]{groupTable, userGroupTable});
                }
                try (ResultSet rs = md.getCrossReference(null, null, groupTable, null, null, userGroupTable)) {
                    if (rs.next()) {
                        groupTableId = rs.getString("PKCOLUMN_NAME");
                        setProperty(USER_TABLE_ID_PARAM,groupTableId);
                        userGroupTableGroupId = rs.getString("FKCOLUMN_NAME");
                        setProperty(USER_GROUP_TABLE_GROUP_ID_PARAM,userGroupTableGroupId);
                    }
                }
            } catch (SQLException e) {
                throw new BadRealmException("An error occurred trying to query the database schema.",e);
            }
        }
    }

    String[] authenticate(String user, char[] password) throws BadRealmException {
        try {
            String userTable = getProperty(USER_TABLE_PARAM);
            String groupTable = getProperty(GROUP_TABLE_PARAM);
            String userGroupTable = getProperty(USER_GROUP_TABLE_PARAM);
            String userNameCol = getProperty(USER_NAME_COLUMN_PARAM);
            String passwordCol = getProperty(PASSWORD_COLUMN_PARAM);
            String saltCol = getProperty(SALT_COLUMN_PARAM);
            String groupNameCol = getProperty(GROUP_NAME_COLUMN_PARAM);
            String getUserGroups;
            int keylen = Integer.parseInt(getProperty(DEFAULT_KEY_LENGTH_PARAM))/8;
            int saltlen = Integer.parseInt(getProperty(DEFAULT_SALT_LENGTH_PARAM))/8;
            byte[] hash = null;
            byte[] salt = null;
            boolean nouser = false;
            String[] groups;

            String getUser = String.format("select %s,%s from %s where %s = ?",passwordCol,saltCol,userTable,userNameCol);
            LOG.log(Level.FINER, "fetch user SQL: {0}", getUser);
            if (userGroupTable == null) {
                LOG.finer("Using simple JDBC Realm-style config");
                getUserGroups = String.format("select %s from %s where %s = ?",groupNameCol,groupTable,userNameCol);
                LOG.log(Level.FINER, "fetch user groups SQL: {0}", getUserGroups);
            } else {
                LOG.finer("Using full-featured config");
                determineUserGroupsSchema();
                String userTableId = getProperty(USER_TABLE_ID_PARAM);
                String groupTableId = getProperty(GROUP_TABLE_ID_PARAM);
                String userGroupTableUserId = getProperty(USER_GROUP_TABLE_USER_ID_PARAM);
                String userGroupTableGroupId = getProperty(USER_GROUP_TABLE_GROUP_ID_PARAM);
                getUserGroups = String.format("select %s.%s from %1$s join %s on %1$s.%s = %3$s.%s join %s on %3$s.%s = %6$s.%s where %6$s.%s = ?",
                        groupTable,groupNameCol,userGroupTable,groupTableId,userGroupTableGroupId,userTable,userGroupTableUserId,userTableId,userNameCol);
                LOG.log(Level.FINER, "fetch user groups SQL: {0}", getUserGroups);
            }
            try (Connection c = getConnection()) {
                LOG.finer("Preparing statements");
                try (PreparedStatement stmt1 = c.prepareStatement(getUser);PreparedStatement stmt2 = c.prepareStatement(getUserGroups)) {
                    stmt1.setString(1, user);
                    stmt2.setString(1, user);
                    LOG.finer("Preparing statements");
                    try (ResultSet users = stmt1.executeQuery();
                            ResultSet userGroups = stmt2.executeQuery()) {
                        Set<String> groupSet = new HashSet<>();
                        if (users.next()) {
                            hash = users.getBytes(1);
                            salt = users.getBytes(2);
                            if (LOG.isLoggable(Level.FINER)) {
                                LOG.log(Level.FINER, "Got hash: {0}; salt: {1}", new Object[]{new BigInteger(1, hash).toString(16), new BigInteger(1,salt).toString(16)});
                            }
                        }
                        while (userGroups.next()) {
                            groupSet.add(userGroups.getString(1));
                        }
                        groups = groupSet.toArray(new String[groupSet.size()]);
                        if (LOG.isLoggable(Level.FINER)) {
                            LOG.log(Level.FINER, "User in groups: {0}", Arrays.toString(groups));
                        }
                    }
                }
            } catch (SQLException e) {
                BadRealmException ee = new BadRealmException(e);
                LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "authenticate", ee);
                throw ee;
            }
            if (hash == null) {
                hash = new byte[keylen];
                nouser = true;
            }
            if (salt == null) {
                salt = new byte[saltlen];
                nouser = true;
            }
            byte[] testhash = hash(password, salt, hash.length * 8);
            if (equalsLengthConstantTime(hash, testhash) && !nouser) {
                return addAssignGroups(groups);
            }
            return null;
        } finally {
            Arrays.fill(password, '\0');
        }
    }

    @Override
    public Enumeration getGroupNames(String string) throws InvalidOperationException, NoSuchUserException {
        String userTable = getProperty(USER_TABLE_PARAM);
        String groupTable = getProperty(GROUP_TABLE_PARAM);
        String userGroupTable = getProperty(USER_GROUP_TABLE_PARAM);
        String userNameCol = getProperty(USER_NAME_COLUMN_PARAM);
        String groupNameCol = getProperty(GROUP_NAME_COLUMN_PARAM);
        String getUserGroups;
        String[] groups = null;
        String hasUser = String.format("select count(*) from %s where %s = ?");
        if (userGroupTable == null) {
            LOG.finer("Using simple JDBC Realm-style config");
            getUserGroups = String.format("select %s from %s where %s = ?",groupNameCol,groupTable,userNameCol);
            LOG.log(Level.FINER, "fetch user groups SQL: {0}", getUserGroups);
        } else {
            LOG.finer("Using full-featured config");
            try {
                determineUserGroupsSchema();
            } catch (BadRealmException e) {
                RuntimeException ee = new RuntimeException(e);
                LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm","getGroupNamed",ee);
                throw ee;
            }
            String userTableId = getProperty(USER_TABLE_ID_PARAM);
            String groupTableId = getProperty(GROUP_TABLE_ID_PARAM);
            String userGroupTableUserId = getProperty(USER_GROUP_TABLE_USER_ID_PARAM);
            String userGroupTableGroupId = getProperty(USER_GROUP_TABLE_GROUP_ID_PARAM);
            getUserGroups = String.format("select %s.%s from %1$s join %s on %1$s.%s = %3$s.%s join %s on %3$s.%s = %6$s.%s where %6$s.%s = ?",
                    groupTable,groupNameCol,userGroupTable,groupTableId,userGroupTableGroupId,userTable,userGroupTableUserId,userTableId,userNameCol);
            LOG.log(Level.FINER, "fetch user groups SQL: {0}", getUserGroups);
        }
        try (Connection c = getConnection()) {
            LOG.finer("Preparing statements");
            try (PreparedStatement stmt1 = c.prepareStatement(getUserGroups); PreparedStatement stmt2 = c.prepareStatement(hasUser)) {
                stmt1.setString(1, string);
                stmt2.setString(1, string);
                LOG.finer("Preparing statements");
                try (ResultSet userGroups = stmt1.executeQuery(); ResultSet userCount = stmt2.executeQuery()) {
                    if (userCount.next()) {
                        if (userCount.getInt(1) == 0) {
                            throw new NoSuchUserException(string);
                        }
                    } else {
                        throw new NoSuchUserException(string);
                    }
                    Set<String> groupSet = new HashSet<>();
                    while (userGroups.next()) {
                        groupSet.add(userGroups.getString(1));
                    }
                    groups = groupSet.toArray(new String[groupSet.size()]);
                    if (LOG.isLoggable(Level.FINER)) {
                        LOG.log(Level.FINER, "User in groups: {0}", Arrays.toString(groups));
                    }
                }
            }
        } catch (SQLException | BadRealmException e) {
            LOG.log(Level.INFO, "A SQL exception occurred. We probably got a BadRealmException earlier, so this is ignored", e);
        }
        groups = addAssignGroups(groups);
        return new Vector<>(Arrays.asList(groups)).elements();
    }

    private byte[] hash(char[] password, byte[] salt, int keylen) throws BadRealmException {
        try {
            LOG.finer("Hashing password");
            int rounds = Integer.parseInt(getProperty(ROUNDS_PARAM));
            PBEKeySpec spec = new PBEKeySpec(password, salt, rounds, keylen);
            Arrays.fill(password, '\0');
            SecretKey key = fac.generateSecret(spec);
            spec.clearPassword();
            return key.getEncoded();
        } catch (InvalidKeySpecException e) {
            throw new BadRealmException("Could not generate password hash",e);
        } finally {
            Arrays.fill(password, '\0');
        }
    }

    private boolean equalsLengthConstantTime(byte[] a, byte[] b) {
        int ct = Math.max(a.length, b.length);
        int rslt = a.length ^ b.length;
        for (int i = 0; i < ct; i++) {
            rslt |= a[Math.min(i,a.length-1)] ^ b[Math.min(i,b.length-1)];
        }
        return rslt == 0;
    }

    private Connection getConnection() throws BadRealmException, SQLException {
        String dbJndi = getProperty(DATASOURCE_JNDI_PARAM);
        String dbUser = getProperty(DB_USER_PARAM);
        String dbPassword = getProperty(DB_PASSWORD_PARAM);
        LOG.log(Level.FINER, "Connecting to {0}", dbJndi);
        try (AutoCloseableContext c = new AutoCloseableContext(new InitialContext())) {
            DataSource ds = (DataSource) c.lookup(dbJndi);
            if (dbUser != null && dbPassword != null) {
                return ds.getConnection(dbUser, dbPassword);
            } else {
                return ds.getConnection();
            }
        } catch (NamingException e) {
            BadRealmException ee = new BadRealmException("Cannot find the JNDI reference to the database",e);
            LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "getConnection", ee);
            throw ee;
        } catch (SQLException e) {
            BadRealmException ee = new BadRealmException("Cannot connect to the database",e);
            LOG.throwing("org.papillion39.sechelper.modules.authentication.JDBCSaltedDigestRealm", "getConnection", ee);
            throw ee;
        }
    }

    public static class AutoCloseableContext implements Context, AutoCloseable {
        private final Context context;

        public AutoCloseableContext(Context context) {
            this.context = context;
        }

        @Override
        public Object lookup(Name name) throws NamingException {
            return context.lookup(name);
        }

        @Override
        public Object lookup(String name) throws NamingException {
            return context.lookup(name);
        }

        @Override
        public void bind(Name name, Object obj) throws NamingException {
            context.bind(name, obj);
        }

        @Override
        public void bind(String name, Object obj) throws NamingException {
            context.bind(name, obj);
        }

        @Override
        public void rebind(Name name, Object obj) throws NamingException {
            context.rebind(name, obj);
        }

        @Override
        public void rebind(String name, Object obj) throws NamingException {
            context.rebind(name, obj);
        }

        @Override
        public void unbind(Name name) throws NamingException {
            context.unbind(name);
        }

        @Override
        public void unbind(String name) throws NamingException {
            context.unbind(name);
        }

        @Override
        public void rename(Name oldName, Name newName) throws NamingException {
            context.rename(oldName, newName);
        }

        @Override
        public void rename(String oldName, String newName) throws NamingException {
            context.rename(oldName, newName);
        }

        @Override
        public NamingEnumeration<NameClassPair> list(Name name) throws NamingException {
            return context.list(name);
        }

        @Override
        public NamingEnumeration<NameClassPair> list(String name) throws NamingException {
            return context.list(name);
        }

        @Override
        public NamingEnumeration<Binding> listBindings(Name name) throws NamingException {
            return context.listBindings(name);
        }

        @Override
        public NamingEnumeration<Binding> listBindings(String name) throws NamingException {
            return context.listBindings(name);
        }

        @Override
        public void destroySubcontext(Name name) throws NamingException {
            context.destroySubcontext(name);
        }

        @Override
        public void destroySubcontext(String name) throws NamingException {
            context.destroySubcontext(name);
        }

        @Override
        public Context createSubcontext(Name name) throws NamingException {
            return context.createSubcontext(name);
        }

        @Override
        public Context createSubcontext(String name) throws NamingException {
            return context.createSubcontext(name);
        }

        @Override
        public Object lookupLink(Name name) throws NamingException {
            return context.lookupLink(name);
        }

        @Override
        public Object lookupLink(String name) throws NamingException {
            return context.lookupLink(name);
        }

        @Override
        public NameParser getNameParser(Name name) throws NamingException {
            return context.getNameParser(name);
        }

        @Override
        public NameParser getNameParser(String name) throws NamingException {
            return context.getNameParser(name);
        }

        @Override
        public Name composeName(Name name, Name prefix) throws NamingException {
            return context.composeName(name, prefix);
        }

        @Override
        public String composeName(String name, String prefix) throws NamingException {
            return context.composeName(name, prefix);
        }

        @Override
        public Object addToEnvironment(String propName, Object propVal) throws NamingException {
            return context.addToEnvironment(propName, propVal);
        }

        @Override
        public Object removeFromEnvironment(String propName) throws NamingException {
            return context.removeFromEnvironment(propName);
        }

        @Override
        public Hashtable<?, ?> getEnvironment() throws NamingException {
            return context.getEnvironment();
        }

        @Override
        public void close() throws NamingException {
            context.close();
        }

        @Override
        public String getNameInNamespace() throws NamingException {
            return context.getNameInNamespace();
        }
        
    }
}
